#include <queue>
#include <vector>
#include <arm_neon.h>
#include <pthread.h>
#include <algorithm>
#include <atomic>
#include <cassert>
#include <thread>

// IVF-PQ 全局结构（线程安全）
struct IVF_PQ_Index {
    std::vector<std::vector<float>> centroids;  // IVF聚类中心
    std::vector<std::vector<uint32_t>> lists;   // 倒排列表
    size_t nlist = 32;                         // 聚类中心数量100
    size_t sub_dim;                             // PQ子空间维度
    size_t m = 4;                               // PQ子空间数
    size_t ks = 64;                            // PQ每子空间簇数256
    std::vector<std::vector<float>> pq_codebooks;
    std::vector<std::vector<uint8_t>> pq_codes;
    std::atomic<bool> initialized{false};
};

// 全局变量
static IVF_PQ_Index ivf_pq_index;
static pthread_once_t ivf_pq_init_once = PTHREAD_ONCE_INIT;
static pthread_mutex_t init_mutex = PTHREAD_MUTEX_INITIALIZER;

// 初始化参数传递结构
struct InitParams {
    const float* base;
    size_t base_number;
    size_t vecdim;
};
static InitParams init_params;

// 线程参数结构
struct ThreadData {
    const float* base;
    const float* query;
    const std::vector<uint32_t>* candidates;
    size_t start_idx;
    size_t end_idx;
    size_t vecdim;
    std::vector<float>* dists;
};



// 初始化函数
void init_ivf_pq(const float* base, size_t base_number, size_t vecdim) {
    pthread_mutex_lock(&init_mutex);
    if (ivf_pq_index.initialized) {
        pthread_mutex_unlock(&init_mutex);
        return;
    }

    ivf_pq_index.sub_dim = vecdim / ivf_pq_index.m;
    ivf_pq_index.centroids.resize(ivf_pq_index.nlist);
    
    // 1. IVF聚类初始化
    std::vector<size_t> indices(base_number);
    std::iota(indices.begin(), indices.end(), 0);
    std::random_shuffle(indices.begin(), indices.end());
    
    for (size_t i = 0; i < ivf_pq_index.nlist; ++i) {
        assert(i < indices.size() && "Cluster center index out of range");
        ivf_pq_index.centroids[i].resize(vecdim);
        std::copy_n(base + indices[i] * vecdim, vecdim, ivf_pq_index.centroids[i].data());
    }

    // 2. 构建倒排列表
    ivf_pq_index.lists.resize(ivf_pq_index.nlist);
    for (size_t i = 0; i < base_number; ++i) {
        float min_dist = std::numeric_limits<float>::max();
        size_t best_cluster = 0;
        const float* vec = base + i * vecdim;
        
        for (size_t j = 0; j < ivf_pq_index.nlist; ++j) {
            float dist = 0.0f;
            const float* c = ivf_pq_index.centroids[j].data();
            
            float32x4_t sum = vdupq_n_f32(0.0f);
            size_t d;
            for (d = 0; d + 4 <= vecdim; d += 4) {
                float32x4_t v = vld1q_f32(vec + d);
                float32x4_t diff = vsubq_f32(v, vld1q_f32(c + d));
                sum = vmlaq_f32(sum, diff, diff);
            }
            float32x2_t sum2 = vadd_f32(vget_low_f32(sum), vget_high_f32(sum));
            dist += vget_lane_f32(sum2, 0) + vget_lane_f32(sum2, 1);
            for (; d < vecdim; ++d) dist += (vec[d] - c[d]) * (vec[d] - c[d]);
            
            if (dist < min_dist) {
                min_dist = dist;
                best_cluster = j;
            }
        }
        ivf_pq_index.lists[best_cluster].push_back(i);
    }

    // 3. PQ编码初始化
    ivf_pq_index.pq_codebooks.resize(ivf_pq_index.m);
    ivf_pq_index.pq_codes.resize(base_number, std::vector<uint8_t>(ivf_pq_index.m));
    for (size_t k = 0; k < ivf_pq_index.m; ++k) {
        std::vector<std::vector<float>> sub_vectors(base_number, std::vector<float>(ivf_pq_index.sub_dim));
        for (size_t i = 0; i < base_number; ++i) {
            std::copy_n(base + i * vecdim + k * ivf_pq_index.sub_dim, ivf_pq_index.sub_dim, sub_vectors[i].data());
        }

        std::vector<std::vector<float>> cluster_centers(ivf_pq_index.ks, std::vector<float>(ivf_pq_index.sub_dim));
        std::vector<size_t> assignments(base_number);

        std::random_shuffle(sub_vectors.begin(), sub_vectors.end());
        for (size_t j = 0; j < ivf_pq_index.ks; ++j) {
            cluster_centers[j] = sub_vectors[j];
        }

        for (size_t iter = 0; iter < 5; ++iter) {          //原来是10
            // Assign points
            for (size_t i = 0; i < base_number; ++i) {
                float min_dist = std::numeric_limits<float>::max();
                size_t best_cluster = 0;
                for (size_t j = 0; j < ivf_pq_index.ks; ++j) {
                    float dist = 0.0f;
                    for (size_t d = 0; d < ivf_pq_index.sub_dim; ++d) {
                        dist += (sub_vectors[i][d] - cluster_centers[j][d]) * (sub_vectors[i][d] - cluster_centers[j][d]);
                    }
                    if (dist < min_dist) {
                        min_dist = dist;
                        best_cluster = j;
                    }
                }
                assignments[i] = best_cluster;
            }

            // Recalculate centers
            std::vector<size_t> counts(ivf_pq_index.ks, 0);
            std::vector<std::vector<float>> sums(ivf_pq_index.ks, std::vector<float>(ivf_pq_index.sub_dim, 0.0f));
            for (size_t i = 0; i < base_number; ++i) {
                counts[assignments[i]]++;
                for (size_t d = 0; d < ivf_pq_index.sub_dim; ++d) {
                    sums[assignments[i]][d] += sub_vectors[i][d];
                }
            }
            for (size_t j = 0; j < ivf_pq_index.ks; ++j) {
                for (size_t d = 0; d < ivf_pq_index.sub_dim; ++d) {
                    cluster_centers[j][d] = sums[j][d] / std::max<size_t>(1, counts[j]);
                }
            }
        }

        // Store codebook
        ivf_pq_index.pq_codebooks[k].resize(ivf_pq_index.ks * ivf_pq_index.sub_dim);
        for (size_t j = 0; j < ivf_pq_index.ks; ++j) {
            std::copy_n(cluster_centers[j].data(), ivf_pq_index.sub_dim, 
                       ivf_pq_index.pq_codebooks[k].data() + j * ivf_pq_index.sub_dim);
        }

        // Encode vectors
        for (size_t i = 0; i < base_number; ++i) {
            float min_dist = std::numeric_limits<float>::max();
            uint8_t best_codeword = 0;
            for (uint8_t j = 0; j < ivf_pq_index.ks; ++j) {
                float dist = 0.0f;
                for (size_t d = 0; d < ivf_pq_index.sub_dim; ++d) {
                    dist += (sub_vectors[i][d] - ivf_pq_index.pq_codebooks[k][j * ivf_pq_index.sub_dim + d]) * 
                          (sub_vectors[i][d] - ivf_pq_index.pq_codebooks[k][j * ivf_pq_index.sub_dim + d]);
                }
                if (dist < min_dist) {
                    min_dist = dist;
                    best_codeword = j;
                }
            }
            ivf_pq_index.pq_codes[i][k] = best_codeword;
        }
    }

    ivf_pq_index.initialized.store(true);
    pthread_mutex_unlock(&init_mutex);
}

// 初始化包装函数
void init_ivf_pq_wrapper() {
    init_ivf_pq(init_params.base, init_params.base_number, init_params.vecdim);
}

// 线程函数：精确距离计算（内积距离）
void* compute_inner_product_distances(void* arg) {
    ThreadData* data = static_cast<ThreadData*>(arg);
    for (size_t i = data->start_idx; i < data->end_idx; ++i) {
        assert(i < data->candidates->size() && "Candidate index out of range");
        uint32_t idx = (*data->candidates)[i];
        const float* vec = data->base + idx * data->vecdim;
        float ip = 0.0f;

        for (size_t d = 0; d < data->vecdim; ++d) {
            ip += vec[d] * data->query[d];
        }

        (*data->dists)[i] = 1.0f - ip;
    }
    return nullptr;
}

// 主查询函数
std::priority_queue<std::pair<float, uint32_t>, std::vector<std::pair<float, uint32_t>>, std::greater<std::pair<float, uint32_t>>> 
pthread_ivf_pq(float* base, float* query, size_t base_number, size_t vecdim, size_t k) 
{
    // 设置初始化参数
    init_params = {base, base_number, vecdim};
    
    // 线程安全初始化
    pthread_once(&ivf_pq_init_once, &init_ivf_pq_wrapper);
    
    // 等待初始化完成
    while (!ivf_pq_index.initialized.load()) {
        std::this_thread::yield();
    }

    // 1. 粗搜索找到最近的nprobe个聚类中心
    const size_t nprobe = 10;
    std::vector<std::pair<float, size_t>> centroid_dists(ivf_pq_index.nlist);
    for (size_t i = 0; i < ivf_pq_index.nlist; ++i) {
        float dist = 0.0f;
        const float* c = ivf_pq_index.centroids[i].data();
        
        // NEON加速距离计算
        float32x4_t sum = vdupq_n_f32(0.0f);
        size_t d;
        for (d = 0; d + 4 <= vecdim; d += 4) {
            float32x4_t q = vld1q_f32(query + d);
            float32x4_t diff = vsubq_f32(q, vld1q_f32(c + d));
            sum = vmlaq_f32(sum, diff, diff);
        }
        float32x2_t sum2 = vadd_f32(vget_low_f32(sum), vget_high_f32(sum));
        dist += vget_lane_f32(sum2, 0) + vget_lane_f32(sum2, 1);
        for (; d < vecdim; ++d) dist += (query[d] - c[d]) * (query[d] - c[d]);
        centroid_dists[i] = {dist, i};
    }
    std::sort(centroid_dists.begin(), centroid_dists.end());

    // 2. 收集候选向量
    std::vector<uint32_t> candidates;
    for (size_t i = 0; i < nprobe; ++i) {
        size_t cluster_id = centroid_dists[i].second;
        candidates.insert(candidates.end(), 
                         ivf_pq_index.lists[cluster_id].begin(),
                         ivf_pq_index.lists[cluster_id].end());
    }

    // 3. 多线程精确计算（内积距离）
    const size_t num_threads = 4;
    std::vector<pthread_t> threads(num_threads);
    std::vector<ThreadData> thread_data(num_threads);
    std::vector<float> dists(candidates.size());
    
    size_t chunk_size = candidates.size() / num_threads;
    for (size_t t = 0; t < num_threads; ++t) {
        thread_data[t] = {
            base, query, &candidates,
            t * chunk_size,
            (t == num_threads-1) ? candidates.size() : (t+1)*chunk_size,
            vecdim,
            &dists
        };
        pthread_create(&threads[t], nullptr, compute_inner_product_distances, &thread_data[t]);
    }
    for (auto& thread : threads) {
        pthread_join(thread, nullptr);
    }

    // 4. 构建优先队列（最小堆）
    std::priority_queue<std::pair<float, uint32_t>, std::vector<std::pair<float, uint32_t>>, 
                       std::greater<std::pair<float, uint32_t>>> q;

    for (size_t i = 0; i < candidates.size(); ++i) {
        if (q.size() < k) {
            q.emplace(dists[i], candidates[i]);
        } else if (dists[i] < q.top().first) {
            q.pop();
            q.emplace(dists[i], candidates[i]);
        }
    }

    return q;
}